// SPDX-FileCopyrightText: 2023-2025 Sangfor Technologies Inc.
// SPDX-License-Identifier: Mulan PSL v2
package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"

	"golang.org/x/crypto/pbkdf2"
)

const (
	MaxKeyLen        = 32
	AESKeyIter       = 100000
	AESKeyLength     = 32
	AESKeyPartLength = 32
)

func NewAesCbcSxf(key1Path string, key2Path string, key3 []byte, saltPath string) (*AesCbc, error) {

	key1, err := ioutil.ReadFile(key1Path)
	if err != nil {
		return nil, err
	}
	key2, err := ioutil.ReadFile(key2Path)
	if err != nil {
		return nil, err
	}
	salt, err := ioutil.ReadFile(saltPath)
	if err != nil {
		return nil, err
	}
	if len(key1) != AESKeyPartLength || len(key2) != AESKeyPartLength || len(key3) != AESKeyPartLength {
		return nil, fmt.Errorf("part lenght is not %v", AESKeyLength)
	}
	password := make([]byte, AESKeyPartLength)
	for i := 0; i < AESKeyPartLength; i++ {
		password[i] = key1[i] ^ key2[i] ^ key3[i]
	}

	dk := pbkdf2.Key(password, salt, AESKeyIter, AESKeyLength, sha256.New)

	return NewAesCbc(string(dk))
}

func NewAesCbc(key string) (*AesCbc, error) {
	if len(key) > MaxKeyLen {
		return nil, fmt.Errorf("key too long, max is %v", MaxKeyLen)
	}
	addBytes := bytes.Repeat([]byte(" "), MaxKeyLen-len(key))
	cipherKey := append([]byte(key), addBytes...)

	block, err := aes.NewCipher(cipherKey)
	if err != nil {
		return nil, err
	}

	return &AesCbc{Key: key, CipherKey: cipherKey, Block: block}, nil
}

// 算法缺陷：变得太长了。如果长度不是BlockSIze的整数倍就补充，如果是就再补充一倍
type AesCbc struct {
	Key       string
	CipherKey []byte
	Block     cipher.Block
}

// 编码: aes-cbc
func (a *AesCbc) Encrypt(raw string) string {
	rawData := a.PKCS7Padding([]byte(raw), aes.BlockSize)
	ciphertext := make([]byte, aes.BlockSize+len(rawData))
	iv := ciphertext[:aes.BlockSize]
	io.ReadFull(rand.Reader, iv)

	mode := cipher.NewCBCEncrypter(a.Block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], rawData)
	return hex.EncodeToString(ciphertext)
}

// 解码: aes-cbc
func (a *AesCbc) Decrypt(enc string) (string, error) {
	ciphertext, err := hex.DecodeString(enc)
	if err != nil {
		return "", err
	}
	if len(ciphertext) < aes.BlockSize {
		return "", fmt.Errorf("ciphertext too short, less then %v", aes.BlockSize)
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]
	if len(ciphertext)%aes.BlockSize != 0 {
		return "", fmt.Errorf("ciphertext is not a multiple of the block size")
	}
	mode := cipher.NewCBCDecrypter(a.Block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)

	return string(a.PKCS7UnPadding(ciphertext)), nil
}

// 扩充字节数组成blockSize的整数倍，扩充字节的内容为 扩充长度的值。 本来是整数倍的时候再加一倍
func (a *AesCbc) PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

// 去掉编码时扩充的字节
func (a *AesCbc) PKCS7UnPadding(origData []byte) []byte {
	length := len(origData)
	// 源数据为空，直接返回
	if length == 0 {
		return origData
	}

	unpadding := int(origData[length-1])
	if length < unpadding {
		return origData
	}

	return origData[:(length - unpadding)]
}
